package com.hunterai.controller;

import com.hunterai.entity.GenerateTask;
import com.hunterai.service.GenerateTaskQueueService;
import com.hunterai.comfyui.ComfyUIService;
import com.hunterai.security.JwtUtil;
import com.hunterai.dto.ReturnData;

import org.postgresql.shaded.com.ongres.scram.common.bouncycastle.pbkdf2.Integers;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.server.ResponseStatusException;

import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/api/queue")
public class GenerateTaskQueueController {
    @Autowired
    private GenerateTaskQueueService queueService;

    @Autowired
    private ComfyUIService comfyUIService;
    
    @Autowired
    private JwtUtil jwtUtil;

    // 获取当前登录用户的ID
    private Long getCurrentUserId() {
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        if (authentication != null && authentication.isAuthenticated()) {
            String username = authentication.getName();
            // 从JWT中获取用户ID (需保证用户已通过JWT认证)
            String token = null;
            // 这里假设你可以从请求上下文中获取token
            // 实际实现时，你可能需要从请求头中获取
            return jwtUtil.extractUserId(token);
        }
        throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "未登录用户");
    }

    @PostMapping(value = "/generate", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public ReturnData<Map<String, String>> generate(
        @RequestParam String workflowId,
        @RequestPart(value = "image", required = false) MultipartFile image,
        @RequestParam Map<String, Object> params,
        @RequestParam(required = false) Integer piority,
        @RequestHeader("Authorization") String authHeader // 添加获取Authorization头
    ) {
        //从Authorization头中提取用户ID

        Long userId = null;
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            String token = authHeader.substring(7);
            userId = jwtUtil.extractUserId(token);
        }

        if (userId == null) {
            throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "未登录用户");
        }
         
        // 1. 文件上传，获得图片名
        if (image != null && !image.isEmpty()) {
            String imageName = comfyUIService.imageUpload(image);
            params.put("image", imageName);
        }
        // 2. 提交任务到队列，并传入用户ID
        String taskId = queueService.submitTask(workflowId, params, piority, userId);
        return ReturnData.success(Map.of("taskId", taskId));
    }

    @GetMapping("/checktaskstatus")
    public ReturnData<GenerateTask> checkTaskStatus(
        @RequestParam String taskId,
        @RequestHeader("Authorization") String authHeader
    ) {
        // 从Authorization头中提取用户ID
        Long userId = null;
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            String token = authHeader.substring(7);
            userId = jwtUtil.extractUserId(token);
        }
        
        if (userId == null) {
            throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "未登录用户");
        }
        
        // 验证任务归属
        GenerateTask task = queueService.getTaskByTaskIdAndUserId(taskId, userId);
        if (task == null) {
            // 尝试使用旧方法找任务（向下兼容）
            task = queueService.getTaskByTaskId(taskId);
            // 如果找到了任务但没有用户ID，或者用户ID匹配，则允许查看
            if (task != null && (task.getUserId() == null || task.getUserId().equals(userId))) {
                return ReturnData.success(task);
            }
            throw new ResponseStatusException(HttpStatus.NOT_FOUND, "任务不存在或无权限查看");
        }
        return ReturnData.success(task);
    }
    
    // 添加获取用户任务列表的接口
    @GetMapping("/user-tasks")
    public List<GenerateTask> getUserTasks(@RequestHeader("Authorization") String authHeader) {
        // 从Authorization头中提取用户ID
        Long userId = null;
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            String token = authHeader.substring(7);
            userId = jwtUtil.extractUserId(token);
        }
        
        if (userId == null) {
            throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "未登录用户");
        }
        
        return queueService.getUserTasks(userId);
    }

    @GetMapping("/user-tasks/history")
    public ReturnData<Page<GenerateTask>> getUserTaskHistory(
            @RequestHeader("Authorization") String authHeader,
            @RequestParam(defaultValue = "0") int page,
            @RequestParam(defaultValue = "10") int size) {
        
        // 从Authorization头中提取用户ID
        Long userId = null;
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            String token = authHeader.substring(7);
            userId = jwtUtil.extractUserId(token);
        }
        
        if (userId == null) {
            throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "未登录用户");
        }
        
        PageRequest pageRequest = PageRequest.of(page, size);
        Page<GenerateTask> tasks = queueService.getUserTasksWithPagination(userId, pageRequest);
        return ReturnData.success(tasks);
    }

    @GetMapping("/user-tasks/history/{workflowId}")
    public Page<GenerateTask> getUserTasksByWorkflow(
            @PathVariable String workflowId,
            @RequestParam(defaultValue = "0") int page,
            @RequestParam(defaultValue = "10") int size,
            @RequestHeader("Authorization") String authHeader) {
        
        // 从Authorization头中提取用户ID
        Long userId = null;
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            String token = authHeader.substring(7);
            userId = jwtUtil.extractUserId(token);
        }
        
        if (userId == null) {
            throw new ResponseStatusException(HttpStatus.UNAUTHORIZED, "未登录用户");
        }
        
        // 创建分页请求
        Pageable pageable = PageRequest.of(page, size, Sort.by("createdAt").descending());
        
        // 获取指定workflow的历史记录
        return queueService.findByUserIdAndWorkflowId(userId, workflowId, pageable);
    }
}